from typing import Optional, Union, List, Dict, Any, Callable, Literal
from bson import ObjectId
from pymilvus import (
    FieldSchema,  DataType,
)
import traceback
from functools import wraps


# ======================================================
#  通用安全执行函数（支持数据库上下文）
# ======================================================
def safe_action(action: str = "操作", default: Any = None):
    """
    通用安全装饰器：
    - 自动输出执行日志
    - 捕获异常并返回默认值
    - 支持 show 控制是否显示细节
    """
    def decorator(func: Callable):
        @wraps(func)
        def wrapper(*args, **kwargs):
            self_obj = args[0] if args else None
            alias = getattr(self_obj, "alias", "unknown")
            name = getattr(self_obj, "name", "")
            show_trace = getattr(self_obj, "show", False)
            prefix = f"[Milvus@{alias}/{name}] " if name else f"[Milvus@{alias}] "

            try:
                result = func(*args, **kwargs)
                if show_trace:
                    print(f"{prefix}{action}完成：{summarize_result(result)}")
                return result
            except Exception as e:
                print(f"{prefix}{action}失败: {type(e).__name__}: {e}")
                if show_trace:
                    traceback.print_exc()
                return default

        return wrapper
    return decorator


def summarize_result(result):
    if result is None:
        return "无返回值"
    if hasattr(result, "insert_count"):
        return f"插入 {result.insert_count} 条"
    if hasattr(result, "delete_count"):
        return f"删除 {result.delete_count} 条"
    if isinstance(result, (list, tuple, set)):
        return f"返回 {len(result)} 条结果"
    if isinstance(result, dict):
        return f"返回字典 {len(result)} 项"
    return str(result)[:200]  # 防止过长打印


@safe_action("启用默认嵌入模型")
def load_default_model(model_embed):
    # 如果用户传入了可调用的向量化函数，直接返回
    if callable(model_embed):
        return model_embed
    # 否则加载默认模型
    from sentence_transformers import SentenceTransformer
    model = SentenceTransformer(
        "sentence-transformers/paraphrase-multilingual-MiniLM-L12-v2"
    )
    return lambda batch: model.encode(batch, show_progress_bar=False)


def normalize_ids(ids: Union[str, ObjectId, List[str], List[ObjectId]]) -> List[str]:
    """统一 id 格式"""
    if isinstance(ids, (str, ObjectId)):
        ids = [ids]
    elif not isinstance(ids, list):
        raise TypeError("参数 ids 必须是 str、ObjectId 或它们的列表")
    return [str(i) for i in ids]


def extract_query_result(res):
    """统一提取 Milvus query 返回的数据"""
    clean_res = []
    for item in res:
        if isinstance(item, dict) and "data" in item:
            clean_res.extend(item["data"])
        else:
            clean_res.append(item)
    return clean_res

# 处理数据库格式


def normalize_fields(fields):
    """
    fields: None | list[tuple] | list[dict]
        推荐格式（描述可省略）：
            [
                ("id", "str", 64),
                ("embedding", "float_vector", 384),
                ("title", "str", 300),
                ("created_at", "int64", 11, "创建时间戳")
            ]
        长度写 None、0、False 都会使用默认值。
    """
    # ===== 默认格式 =====
    if fields is None:
        return [
            FieldSchema(name="id", dtype=DataType.VARCHAR,
                        is_primary=True, max_length=64),
            FieldSchema(name="embedding",
                        dtype=DataType.FLOAT_VECTOR, dim=384),
        ]

    if not isinstance(fields, list):
        raise TypeError("fields 必须是 list 类型（list[tuple] 或 list[dict]）")

    # ===== 类型映射（集中定义） =====
    TYPE_MAP = {
        "str": DataType.VARCHAR,
        "string": DataType.VARCHAR,
        "varchar": DataType.VARCHAR,
        "int": DataType.INT64,
        "int64": DataType.INT64,
        "integer": DataType.INT64,
    }

    # ===== 统一提取输入 =====
    normalized = []
    for f in fields:
        if isinstance(f, dict):
            name, ftype, max_len, desc = (
                f.get("name"), f.get("type"), f.get(
                    "max_length"), f.get("description", "")
            )
        elif isinstance(f, (tuple, list)):
            name, ftype, max_len, desc = (
                (f + (None,) * 4)[:4]  # 自动补齐缺少的参数
            )
        else:
            raise TypeError("fields 中的每一项必须是 dict 或 tuple")

        if not name or not ftype:
            raise ValueError(f"字段定义不完整：{f}")

        normalized.append({
            "name": str(name).strip(),
            "type": str(ftype).lower().strip(),
            "max_length": int(max_len) if max_len else 0,
            "description": str(desc or "").strip(),
        })

    # ===== 校验 id 与 embedding =====
    names = [f["name"] for f in normalized]
    if "id" not in names:
        raise ValueError("缺少必要字段：'id'")
    if "embedding" not in names:
        raise ValueError("缺少必要字段：'embedding'")
    if len(names) != len(set(names)):
        raise ValueError("字段名存在重复")

    def default_len(value, default):
        return default if not value or value < 1 else value

    # ===== 构造 FieldSchema 列表 =====
    field_schemas = []
    for f in normalized:
        name, ftype, max_len, desc = f["name"], f["type"], f["max_length"], f["description"]

        # ---- id 字段 ----
        if name == "id":
            max_len = default_len(max_len, 64)
            field_schemas.append(FieldSchema(
                name="id",
                dtype=DataType.VARCHAR,
                is_primary=True,
                max_length=max_len,
                description=desc
            ))
            continue

        # ---- embedding 字段 ----
        if name == "embedding":
            max_len = default_len(max_len, 384)
            field_schemas.append(FieldSchema(
                name="embedding",
                dtype=DataType.FLOAT_VECTOR,
                dim=max_len,
                description=desc
            ))
            continue

        # ---- 其他字段 ----
        if ftype not in TYPE_MAP:
            raise ValueError(f"不支持的字段类型：{ftype}（仅支持 str 和 int）")

        dtype = TYPE_MAP[ftype]
        if dtype == DataType.VARCHAR:
            max_len = default_len(max_len, 252)
            field_schemas.append(FieldSchema(
                name=name,
                dtype=dtype,
                max_length=max_len,
                description=desc
            ))
        else:
            field_schemas.append(FieldSchema(
                name=name,
                dtype=dtype,
                description=desc
            ))

    return field_schemas

# 处理索引参数


def normalize_index(metric_type: Literal["COSINE", "L2"] = "COSINE"):
    """
    metric_type: 相似度度量方式
        - "COSINE" 适用于文本、语义搜索（默认）
        - "L2" 适用于未归一化的向量，如图像或音频特征
    """
    return {
        "metric_type": metric_type,
        "index_type": "AUTOINDEX",  # 自动选择最优索引类型
        "params": {}  # AUTOINDEX 不需要额外参数
    }

# dtype 映射表（更易读）
dtype_map = {
    0: "NONE",
    1: "BOOL",
    2: "INT8",
    3: "INT16",
    4: "INT32",
    5: "INT64",
    10: "FLOAT",
    11: "DOUBLE",
    20: "VARCHAR",
    21: "STRING",
    100: "BINARY_VECTOR",
    101: "FLOAT_VECTOR",
}
